package algorithm

import (
	"encoding/json"
	"math"
	"math/rand"
	"os"
	"time"

	"b2c-delivery-optimization/internal/model"
)

// LSTMPredictor LSTM预测器
// 为什么要使用LSTM?
// 1. LSTM 能够分析长期的历史信息
// 2. LSTM 能够分析周期性变化
// 3. LSTM 能够预测未来订单趋势
type LSTMPredictor struct {
	// LSTM网络参数
	hiddenSize      int     // 隐藏层大小
	sequenceLength  int     // 输入序列长度
	predictionSteps int     // 预测步长
	learningRate    float64 // 学习率
	// 网络权重
	weights map[string][]float64
	// 训练数据
	trainingData []float64
	// 训练状态
	epochs      int                  // 训练轮数
	loss        float64              // 当前损失
	bestLoss    float64              // 最佳损失
	bestWeights map[string][]float64 // 最佳权重
}

// NewLSTMPredictor 创建新的LSTM预测器
func NewLSTMPredictor(hiddenSize, sequenceLength, predictionSteps int) *LSTMPredictor {
	predictor := &LSTMPredictor{
		hiddenSize:      hiddenSize,
		sequenceLength:  sequenceLength,
		predictionSteps: predictionSteps,
		learningRate:    0.01,
		weights:         make(map[string][]float64),
		bestWeights:     make(map[string][]float64),
	}
	// 初始化权重
	predictor.initializeWeights()
	return predictor
}

// initializeWeights 初始化LSTM的权重
func (p *LSTMPredictor) initializeWeights() {
	// 输入门权重
	p.weights["Wi"] = make([]float64, p.hiddenSize)
	p.weights["Ui"] = make([]float64, p.hiddenSize)
	p.weights["bi"] = make([]float64, p.hiddenSize)
	// 遗忘门权重
	p.weights["Wf"] = make([]float64, p.hiddenSize)
	p.weights["Uf"] = make([]float64, p.hiddenSize)
	p.weights["bf"] = make([]float64, p.hiddenSize)
	// 输出门权重
	p.weights["Wo"] = make([]float64, p.hiddenSize)
	p.weights["Uo"] = make([]float64, p.hiddenSize)
	p.weights["bo"] = make([]float64, p.hiddenSize)
	// 单元状态权重
	p.weights["Wc"] = make([]float64, p.hiddenSize)
	p.weights["Uc"] = make([]float64, p.hiddenSize)
	p.weights["bc"] = make([]float64, p.hiddenSize)
	// 输出层权重
	p.weights["Wy"] = make([]float64, p.hiddenSize)
	p.weights["by"] = make([]float64, 1)
	// 初始化权重为小的随机值
	for i := 0; i < p.hiddenSize; i++ {
		p.weights["Wi"][i] = (rand.Float64()*2 - 1) * 0.1
		p.weights["Ui"][i] = (rand.Float64()*2 - 1) * 0.1
		p.weights["bi"][i] = 0.0
		p.weights["Wf"][i] = (rand.Float64()*2 - 1) * 0.1
		p.weights["Uf"][i] = (rand.Float64()*2 - 1) * 0.1
		p.weights["bf"][i] = 0.0
		p.weights["Wo"][i] = (rand.Float64()*2 - 1) * 0.1
		p.weights["Uo"][i] = (rand.Float64()*2 - 1) * 0.1
		p.weights["bo"][i] = 0.0
		p.weights["Wc"][i] = (rand.Float64()*2 - 1) * 0.1
		p.weights["Uc"][i] = (rand.Float64()*2 - 1) * 0.1
		p.weights["bc"][i] = 0.0
		p.weights["Wy"][i] = (rand.Float64()*2 - 1) * 0.1
	}
	p.weights["by"][0] = 0.0
}

// Predict 预测指定时间的订单数量
func (p *LSTMPredictor) Predict(t time.Time) float64 {
	// 获取最近的输入序列
	sequence := p.getRecentSequence(t)
	// 使用训练好的模型进行预测
	outputs, _ := p.forward(sequence)
	// 返回最后一个预测值
	return outputs[len(outputs)-1]
}

// sigmoid 激活函数
func sigmoid(x float64) float64 {
	return 1.0 / (1.0 + math.Exp(-x))
}

// tanh 激活函数
func tanh(x float64) float64 {
	return math.Tanh(x)
}

// Train 训练LSTM模型
func (p *LSTMPredictor) Train(historicalOrders []*model.Order, epochs int) error {
	// 准备训练数据
	p.prepareTrainingData(historicalOrders)
	// 初始化最佳损失
	p.bestLoss = math.MaxFloat64
	// 训练网络
	for epoch := 0; epoch < epochs; epoch++ {
		p.epochs = epoch
		// 使用BPTT算法训练
		loss := p.trainNetwork()
		// 更新最佳权重
		if loss < p.bestLoss {
			p.bestLoss = loss
			p.bestWeights = make(map[string][]float64)
			for k, v := range p.weights {
				p.bestWeights[k] = make([]float64, len(v))
				copy(p.bestWeights[k], v)
			}
		}
		// 打印训练进度
		if epoch%10 == 0 {
			println("Epoch:", epoch, "Loss:", loss)
		}
	}
	// 使用最佳权重
	p.weights = p.bestWeights
	return nil
}

// prepareTrainingData 准备训练数据
func (p *LSTMPredictor) prepareTrainingData(orders []*model.Order) {
	// 1. 按时间窗口统计订单量
	orderCounts := make(map[time.Time]int)
	for _, order := range orders {
		// 将时间按小时对齐
		timeKey := order.CreateTime.Truncate(time.Hour)
		orderCounts[timeKey]++
	}
	// 2. 获取时间范围
	var minTime, maxTime time.Time
	first := true
	for t := range orderCounts {
		if first {
			minTime = t
			maxTime = t
			first = false
		} else {
			if t.Before(minTime) {
				minTime = t
			}
			if t.After(maxTime) {
				maxTime = t
			}
		}
	}
	// 3. 创建连续的时间序列
	timeSeries := make([]float64, 0)
	features := make([][]float64, 0) // 存储时间特征
	currentTime := minTime
	for !currentTime.After(maxTime) {
		// 获取该时间点的订单量
		count := float64(orderCounts[currentTime])
		// 添加时间特征
		hour := float64(currentTime.Hour()) / 24.0      // 归一化到[0,1]
		weekday := float64(currentTime.Weekday()) / 7.0 // 归一化到[0,1]
		dayOfMonth := float64(currentTime.Day()) / 31.0 // 归一化到[0,1]
		// 组合特征
		feature := []float64{count, hour, weekday, dayOfMonth}
		features = append(features, feature)
		// 添加到时间序列
		timeSeries = append(timeSeries, count)
		// 移动到下一个小时
		currentTime = currentTime.Add(time.Hour)
	}
	// 4. 数据归一化
	p.normalizeData(timeSeries)
	// 5. 保存处理后的数据
	p.trainingData = timeSeries
}

// normalizeData 数据归一化
func (p *LSTMPredictor) normalizeData(data []float64) {
	if len(data) == 0 {
		return
	}
	// 找到最大值和最小值
	minVal := data[0]
	maxVal := data[0]
	for _, val := range data {
		if val < minVal {
			minVal = val
		}
		if val > maxVal {
			maxVal = val
		}
	}
	// 归一化到[0,1]范围
	rangeVal := maxVal - minVal
	if rangeVal == 0 {
		rangeVal = 1 // 避免除以0
	}
	for i := range data {
		data[i] = (data[i] - minVal) / rangeVal
	}
}

// trainNetwork 使用BPTT算法训练网络
func (p *LSTMPredictor) trainNetwork() float64 {
	totalLoss := 0.0
	sequenceCount := len(p.trainingData) - p.sequenceLength - p.predictionSteps
	for i := 0; i < sequenceCount; i++ {
		// 获取输入序列和目标序列
		inputSequence := p.trainingData[i : i+p.sequenceLength]
		targetSequence := p.trainingData[i+p.sequenceLength : i+p.sequenceLength+p.predictionSteps]
		// 前向传播
		outputs, states := p.forward(inputSequence)
		// 计算损失
		loss := p.calculateLoss(outputs, targetSequence)
		totalLoss += loss
		// 反向传播
		p.backward(inputSequence, targetSequence, outputs, states)
	}
	return totalLoss / float64(sequenceCount)
}

// forward 前向传播
func (p *LSTMPredictor) forward(inputs []float64) ([]float64, [][]float64) {
	// 初始化状态
	// 隐藏状态
	h := make([]float64, p.hiddenSize)
	// 记忆状态
	c := make([]float64, p.hiddenSize)
	states := make([][]float64, len(inputs))
	outputs := make([]float64, len(inputs))
	for t := 0; t < len(inputs); t++ {
		// 计算输入门
		i := sigmoid(dot(p.weights["Wi"], []float64{inputs[t]}) + dot(p.weights["Ui"], h) + p.weights["bi"][0])
		iVec := make([]float64, p.hiddenSize)
		for j := range iVec {
			iVec[j] = i
		}
		// 计算遗忘门
		f := sigmoid(dot(p.weights["Wf"], []float64{inputs[t]}) + dot(p.weights["Uf"], h) + p.weights["bf"][0])
		fVec := make([]float64, p.hiddenSize)
		for j := range fVec {
			fVec[j] = f
		}
		// 计算输出门
		o := sigmoid(dot(p.weights["Wo"], []float64{inputs[t]}) + dot(p.weights["Uo"], h) + p.weights["bo"][0])
		oVec := make([]float64, p.hiddenSize)
		for j := range oVec {
			oVec[j] = o
		}
		// 计算候选状态
		c_tilde := tanh(dot(p.weights["Wc"], []float64{inputs[t]}) + dot(p.weights["Uc"], h) + p.weights["bc"][0])
		c_tildeVec := make([]float64, p.hiddenSize)
		for j := range c_tildeVec {
			c_tildeVec[j] = c_tilde
		}
		// 更新单元状态
		c = add(multiply(fVec, c), multiply(iVec, c_tildeVec))
		// 更新隐藏状态
		tanhC := make([]float64, p.hiddenSize)
		for j := range tanhC {
			tanhC[j] = tanh(c[j])
		}
		h = multiply(oVec, tanhC)
		// 计算输出
		outputs[t] = dot(p.weights["Wy"], h) + p.weights["by"][0]
		// 保存状态
		states[t] = make([]float64, p.hiddenSize*2)
		copy(states[t][:p.hiddenSize], h)
		copy(states[t][p.hiddenSize:], c)
	}
	return outputs, states
}

// backward 反向传播
func (p *LSTMPredictor) backward(inputs, targets, outputs []float64, states [][]float64) {
	// 初始化梯度
	dWi := make([]float64, p.hiddenSize)
	dWf := make([]float64, p.hiddenSize)
	dWo := make([]float64, p.hiddenSize)
	dWc := make([]float64, p.hiddenSize)
	dUi := make([]float64, p.hiddenSize)
	dUf := make([]float64, p.hiddenSize)
	dUo := make([]float64, p.hiddenSize)
	dUc := make([]float64, p.hiddenSize)
	dbi := make([]float64, p.hiddenSize)
	dbf := make([]float64, p.hiddenSize)
	dbo := make([]float64, p.hiddenSize)
	dbc := make([]float64, p.hiddenSize)
	dWy := make([]float64, p.hiddenSize)
	dby := make([]float64, 1)
	// 计算输出层梯度
	dy := make([]float64, len(outputs))
	for t := 0; t < len(outputs); t++ {
		dy[t] = outputs[t] - targets[t]
		dyVec := make([]float64, p.hiddenSize)
		for j := range dyVec {
			dyVec[j] = dy[t]
		}
		dWy = add(dWy, multiply(dyVec, states[t][:p.hiddenSize]))
		dby[0] += dy[t]
	}
	// 更新权重
	for i := 0; i < p.hiddenSize; i++ {
		p.weights["Wi"][i] -= p.learningRate * dWi[i]
		p.weights["Wf"][i] -= p.learningRate * dWf[i]
		p.weights["Wo"][i] -= p.learningRate * dWo[i]
		p.weights["Wc"][i] -= p.learningRate * dWc[i]
		p.weights["Ui"][i] -= p.learningRate * dUi[i]
		p.weights["Uf"][i] -= p.learningRate * dUf[i]
		p.weights["Uo"][i] -= p.learningRate * dUo[i]
		p.weights["Uc"][i] -= p.learningRate * dUc[i]
		p.weights["bi"][i] -= p.learningRate * dbi[i]
		p.weights["bf"][i] -= p.learningRate * dbf[i]
		p.weights["bo"][i] -= p.learningRate * dbo[i]
		p.weights["bc"][i] -= p.learningRate * dbc[i]
		p.weights["Wy"][i] -= p.learningRate * dWy[i]
	}
	p.weights["by"][0] -= p.learningRate * dby[0]
}

// calculateLoss 计算损失
func (p *LSTMPredictor) calculateLoss(outputs, targets []float64) float64 {
	loss := 0.0
	for i := 0; i < len(outputs); i++ {
		diff := outputs[i] - targets[i]
		loss += diff * diff
	}
	return loss / float64(len(outputs))
}

// getRecentSequence 获取最近的输入序列
func (p *LSTMPredictor) getRecentSequence(t time.Time) []float64 {
	// 这里简化实现，实际应用中需要从历史数据中获取
	sequence := make([]float64, p.sequenceLength)
	for i := 0; i < p.sequenceLength; i++ {
		sequence[i] = 50.0 * p.getTimeMultiplier(t.Hour())
	}
	return sequence
}

// SaveModel 保存模型
func (p *LSTMPredictor) SaveModel(filename string) error {
	data, err := json.Marshal(p.weights)
	if err != nil {
		return err
	}
	return os.WriteFile(filename, data, 0644)
}

// LoadModel 加载模型
func (p *LSTMPredictor) LoadModel(filename string) error {
	data, err := os.ReadFile(filename)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, &p.weights)
}

// 辅助函数
func dot(a, b []float64) float64 {
	sum := 0.0
	for i := 0; i < len(a); i++ {
		sum += a[i] * b[i]
	}
	return sum
}

func add(a, b []float64) []float64 {
	result := make([]float64, len(a))
	for i := 0; i < len(a); i++ {
		result[i] = a[i] + b[i]
	}
	return result
}

func multiply(a, b []float64) []float64 {
	result := make([]float64, len(a))
	for i := 0; i < len(a); i++ {
		result[i] = a[i] * b[i]
	}
	return result
}

// getTimeMultiplier 根据时间获取订单量倍数
func (p *LSTMPredictor) getTimeMultiplier(hour int) float64 {
	// 早高峰 (7-9点)
	if hour >= 7 && hour <= 9 {
		return 1.5
	}
	// 午高峰 (11-13点)
	if hour >= 11 && hour <= 13 {
		return 2.0
	}
	// 晚高峰 (17-19点)
	if hour >= 17 && hour <= 19 {
		return 2.5
	}
	// 深夜 (23-5点)
	if hour >= 23 || hour <= 5 {
		return 0.5
	}
	// 其他时间
	return 1.0
}
